import os
import json
import glob
import random
import argparse
import shutil

import numpy as np
from collections import defaultdict
from datetime import datetime
import matplotlib.pyplot as plt
import matplotlib.patches as patches

from PIL import Image


# 默认类别名（第一个是 background）
# DEFAULT_CLASS_NAMES = ['__background__', 'person', 'car']


def visualize_annotations(image_path, labelme_data, class_id_to_name):
    """可视化标注框"""
    image = np.array(Image.open(image_path))
    fig, ax = plt.subplots(1)
    ax.imshow(image)

    for shape in labelme_data["shapes"]:
        label = shape["label"]
        if label not in class_id_to_name.values():
            continue

        points = shape["points"]
        x_coords = [p[0] for p in points]
        y_coords = [p[1] for p in points]
        x_min, y_min = min(x_coords), min(y_coords)
        width = max(x_coords) - x_min
        height = max(y_coords) - y_min

        rect = patches.Rectangle((x_min, y_min), width, height,
                                 linewidth=2, edgecolor='r', facecolor='none')
        ax.add_patch(rect)
        ax.text(x_min, y_min, label, color='yellow', fontsize=12,
                bbox=dict(facecolor='black', alpha=0.5))

    plt.axis('off')
    plt.show()


def labelme2coco(labelme_dir, image_dir, output_dir, class_names, train_ratio=0.7, val_ratio=0.2, test_ratio=0.1, visualize=False):
    """
    将 labelme 的 json 文件转换为 COCO 格式，并划分 train/val/test
    :param labelme_dir: 存放 labelme json 文件的目录
    :param image_dir: 存放图像文件的目录
    :param output_dir: 输出 JSON 文件目录
    :param class_names: 类别列表（第一个是 background）
    :param train_ratio: 训练集比例
    :param val_ratio: 验证集比例
    :param test_ratio: 测试集比例
    :param visualize: 是否可视化标注框
    """
    # assert train_ratio + val_ratio + test_ratio == 1.0, "Train/Val/Test 比例总和必须为 1.0"

    class_name_to_id = {name: i for i, name in enumerate(class_names)}
    categories = [{"id": v, "name": k} for k, v in class_name_to_id.items() if k != "__background__"]

    # 统计类别分布
    class_distribution = defaultdict(int)

    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)



    # 获取所有 JSON 文件
    # json_paths = glob.glob(os.path.join(labelme_dir, "*.json"))
    json_paths = glob.glob(labelme_dir)
    random.shuffle(json_paths)

    # total = len(json_paths)
    # train_split = int(total * train_ratio)
    # val_split = int(total * (train_ratio + val_ratio))
    test_split = int(len(json_paths)*0.1)

    splits = {
        "train2017": json_paths,
        # "val": json_paths[train_split:val_split],
        "val2017": json_paths[:test_split]
    }

    os.makedirs(f'{output_dir}/train2017')
    os.makedirs(f'{output_dir}/val2017') # 图片存放

    annotations_save_dir = f'{output_dir}/annotations'
    os.makedirs(annotations_save_dir)
    for split_name, files in splits.items():
        print(f"Processing {split_name} set ({len(files)} files)...")

        coco_dataset = {
            "info": {
                "description": f"{split_name} set converted by labelme2coco",
                "version": "1.0",
                "year": datetime.now().year,
                "contributor": "User",
                "date_created": datetime.now().strftime("%Y/%m/%d")
            },
            "licenses": [],
            "images": [],
            "annotations": [],
            "categories": categories
        }

        ann_id = 1
        for json_path in files:
            with open(json_path, encoding="utf-8") as f:
                labelme_data = json.load(f)

            # image_name = labelme_data["imagePath"]
            # image_path = os.path.join(image_dir, image_name)
            image_path = json_path.replace(".json", ".jpg")
            image_name = os.path.basename(image_path)

            shutil.copy(image_path, f'{output_dir}/{split_name}')
            if not os.path.exists(image_path):
                print(f"Image file not found: {image_path}")
                raise Exception(f"Image file not found: {image_path}")
                # continue


            if visualize:
                visualize_annotations(image_path, labelme_data, {v: k for k, v in class_name_to_id.items()})

            image = Image.open(image_path)

            img_w, img_h = image.size

            image_info = {
                "id": len(coco_dataset["images"]) + 1,
                "file_name": image_name,
                "width": img_w,
                "height": img_h,
                "date_captured": "",
                "license": 0
            }
            coco_dataset["images"].append(image_info)

            for shape in labelme_data.get("shapes", []):
                label = shape["label"]
                if label not in class_name_to_id:
                    print(f"Class '{label}' not in class_names, skipping...")
                    continue
                category_id = class_name_to_id[label]
                class_distribution[label] += 1

                points = shape["points"]
                segmentation = [np.array(points).flatten().tolist()]

                x_coords = [p[0] for p in points]
                y_coords = [p[1] for p in points]
                x_min, y_min = min(x_coords), min(y_coords)
                x_max, y_max = max(x_coords), max(y_coords)
                width = x_max - x_min
                height = y_max - y_min
                bbox = [x_min, y_min, width, height]
                area = width * height

                annotation = {
                    "id": ann_id,
                    "image_id": image_info["id"],
                    "category_id": category_id,
                    "bbox": bbox,
                    "area": area,
                    "segmentation": segmentation,
                    "iscrowd": 0
                }
                coco_dataset["annotations"].append(annotation)
                ann_id += 1

        # 保存 COCO JSON
        output_file = os.path.join(annotations_save_dir, f"instances_{split_name}.json")
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(coco_dataset, f, indent=4)

        print(f"Saved COCO format JSON to {output_file}")

    # 打印类别分布
    print("\n📊 Class Distribution:")
    for cls, count in class_distribution.items():
        print(f"{cls}: {count}")


def main():

    class_names = ['aircraft', 'cargo_door', 'chock', 'refueling_truck', 'lifting_platform',
                   'transfer_car', 'staff', 'connecting', 'catering_truck', 'belt_loader',
                   'power_cable', 'pre_conditioning_air', 'thumbs_up', 'ULD', 'bag', 'push_back_tractor',]
    # output_dir = r"D:\DATA\20250611HKBZ\trainDetr_all\0613_testdata_labelme20_zxc_coco"
    # labelme_dir = r'D:\DATA\20250611HKBZ\trainDetr_all\0613_testdata_labelme20_zxc\*.json'
    # image_dir = r'D:\DATA\20250611HKBZ\trainDetr_all\0613_testdata_labelme20_zxc\*.jpg'


    labelme_dir = r'D:\DATA\20250611HKBZ\trainV8Seg_allv2\add_imgs\*\*.json'
    image_dir = r'D:\DATA\20250611HKBZ\trainV8Seg_allv2\add_imgs\*\*.jpg'
    output_dir = r"D:\DATA\20250611HKBZ\trainDetr_all\coco800_format_data"
    parser = argparse.ArgumentParser(description="Convert Labelme annotations to COCO format with train/val/test split.")
    parser.add_argument("--labelme_dir", type=str, default=labelme_dir, help="Directory containing labelme json files.")
    parser.add_argument("--image_dir", type=str, default=image_dir, help="Directory containing image files.")
    parser.add_argument("--output_dir", type=str, default=output_dir, help="Output directory for COCO JSON files.")
    parser.add_argument("--class_names", nargs='+', default=class_names, help="List of class names (first is background).")
    # parser.add_argument("--train_ratio", type=float, default=1, help="Ratio of training data (e.g., 0.7)")
    # parser.add_argument("--val_ratio", type=float, default=0.0, help="Ratio of validation data (e.g., 0.2)")
    # parser.add_argument("--test_ratio", type=float, default=0.2, help="Ratio of test data (e.g., 0.1)")
    parser.add_argument("--visualize", action="store_true", help="Visualize the annotations")

    args = parser.parse_args()

    labelme2coco(
        labelme_dir=args.labelme_dir,
        image_dir=args.image_dir,
        output_dir=args.output_dir,
        class_names=args.class_names,
        # train_ratio=args.train_ratio,
        # val_ratio=args.val_ratio,
        # test_ratio=args.test_ratio,
        visualize=args.visualize
    )


if __name__ == "__main__":
    main()